import json
import time

from matplotlib import pyplot as plt
from train_mlp_ctrl import NormalizedEndJointData

from model import DiffusionIKModel
from robots.piper import PiperArm

filename = "data/json_data/combined.json"

# with open(filename, "r") as f:
#     data = json.load(f)

data = NormalizedEndJointData(filename)

plt.hist(data.joints[:, 0])


arm = PiperArm("can0", DiffusionIKModel(in_dim=3, out_dim=4, model_file="diffusion_end_joint_v0.ckp"))


for i in range(10):
    arm.move_to_joint(data.joint[1000 * i].tolist())
    time.sleep(0.1)
    print("========================")
    print("real: ", arm.get_arm_pos_status()["position"])
    print("data: ", data.arm[1000 * i])
    print("gripper real:", arm.get_effector_xquat())
    print("gripper data: ", data.gripper[1000 * i])
    print("========================")
    time.sleep(1)
    arm.move_to_arm_pos(data.arm[1000 * i].tolist())
    arm.go_zero()
